import pandas as pd
from nlpx.model import TextCNN
from sklearn.model_selection import train_test_split
from nlpx.llm import AlbertTokenizeVec, ErnieTokenizeVec, train_test_set
from nlpx.model.wrapper import ClassifyModelWrapper

pretrained_path = r'/Users/summy/data/albert_small_zh'

if __name__ == '__main__':
	df = pd.read_csv('~/project/python/parttime/text_gcn/data/北方地区不安全事件统计20240331.csv', encoding='GBK')
	df['故障标志'] = df['故障标志'].astype('category')
	classes = df['故障标志'].cat.categories.values.tolist()
	df['故障标志'] = df['故障标志'].cat.codes
	y = df['故障标志'].to_numpy()

	X_train, X_test, y_train, y_test = train_test_split(df['故障描述'].values, y, test_size=0.2)

	tokenize_vec = AlbertTokenizeVec(pretrained_path)
	train_set, test_set = train_test_set(tokenize_vec, X_train, X_test, y_train, y_test)

	text_cnn = TextCNN(tokenize_vec.hidden_size, out_features=len(classes))
	model = ClassifyModelWrapper(classes=classes)
	model.train(text_cnn, train_set, test_set, early_stopping_rounds=2)

	test_texts = ['昆明航B737飞机执行昆明至西安航班 西安机场进近过程中  机组将飞机襟翼位置设置错误触发近地警告事件',
				  '海航 B737飞机执行汉中至广州航班 飞机起飞后发现汉中过站期间货舱有730KG货物未卸机']
	inputs = tokenize_vec.encode_plus(test_texts)
	print(model.predict_classes_proba(inputs))
